#include "AlgorithmFaceDetection.h"
#include "Config.h"
#include "Utils/Log.h"
#include "Utils/Common.h"
#include "facedetection/facedetectcnn.h"
#include "Utils/AlgUtils.h"

namespace AVSAnalyzer
{
    // 定义一些常量
    #define DETECT_BUFFER_SIZE 0x20000

    FaceDetectionEngine::FaceDetectionEngine(const AlgorithmConfig& algorithmConfig)
        : mAlgorithmConfig(algorithmConfig),
          mScoreThreshold(algorithmConfig.scoreThreshold),
          mFaceDetector(nullptr)
    {
        try {
            // 初始化人脸检测器
            mFaceDetector = new unsigned char[DETECT_BUFFER_SIZE];
            if(!mFaceDetector) {
                throw std::runtime_error("Failed to allocate face detection buffer");
            }

            LOGI("Face detection engine initialized successfully");
        }
        catch (const std::exception& e) {
            LOGE("Failed to initialize face detection engine: %s", e.what());
            throw;
        }
    }

    FaceDetectionEngine::~FaceDetectionEngine()
    {
        if(mFaceDetector) {
            delete[] static_cast<unsigned char*>(mFaceDetector);
            mFaceDetector = nullptr;
        }
    }

    bool FaceDetectionEngine::runInference(cv::Mat& image, std::vector<DetectObject>& detects)
    {
        try {
            // 确保图像是BGR格式
            if(image.empty()) {
                LOGE("Input image is empty");
                return false;
            }

            // 执行人脸检测
            int * pResults = facedetect_cnn(
                static_cast<unsigned char*>(mFaceDetector),
                (unsigned char*)(image.data),
                image.cols,
                image.rows,
                (int)image.step
            );

            // 处理检测结果
            detects.clear();
            if(pResults) {
                int face_count = pResults[0];  // 获取检测到的人脸数量
                
                for(int i = 0; i < face_count; i++) {
                    short * p = ((short*)(pResults + 1)) + 142 * i;
                    int confidence = p[0];
                    int x = p[1];
                    int y = p[2];
                    int w = p[3];
                    int h = p[4];
                    
                    float score = confidence / 100.0f;  // 转换置信度到0-1范围
                    
                    if(score >= mScoreThreshold) {
                        DetectObject obj;
                        obj.x1 = x;
                        obj.y1 = y;
                        obj.x2 = x + w;
                        obj.y2 = y + h;
                        obj.class_score = score;
                        obj.class_id = 0;  // 人脸类别ID设为0
                        obj.class_name = "face";  // 类别名称设为face
                        
                        // 确保坐标在图像范围内
                        obj.x1 = std::max(0, std::min(obj.x1, image.cols));
                        obj.y1 = std::max(0, std::min(obj.y1, image.rows));
                        obj.x2 = std::min(image.cols, obj.x2);
                        obj.y2 = std::min(image.rows, obj.y2);
                        
                        if(obj.x2 > obj.x1 && obj.y2 > obj.y1) {
                            detects.push_back(obj);
                        }
                    }
                }
            }

            return true;
        }
        catch(const std::exception& e) {
            LOGE("Face detection error: %s", e.what());
            return false;
        }
    }

    // 实现 AlgorithmFaceDetection 类的方法
    AlgorithmFaceDetection::AlgorithmFaceDetection(const AlgorithmConfig& algorithmConfig)
        : Algorithm(algorithmConfig)
    {
        mEngine = new FaceDetectionEngine(algorithmConfig);
    }

    AlgorithmFaceDetection::~AlgorithmFaceDetection()
    {
        if(mEngine) {
            delete mEngine;
            mEngine = nullptr;
        }
    }

    bool AlgorithmFaceDetection::objectDetect(cv::Mat& image, std::vector<DetectObject>& detects)
    {
        std::lock_guard<std::mutex> lock(mMtx);
        return mEngine->runInference(image, detects);
    }
}
